/*
 * MIT License
 *
 * Copyright (c) 2020 wen.gu <454727014@qq.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

/***************************************************************************
* Name: av_demux.cpp
*
* Purpose: demux component implementation with libav
*
* Developer:
*   wen.gu , 2018-12-29
*
* TODO:
*
***************************************************************************/

/******************************************************************************
**    INCLUDES
******************************************************************************/
#define _CRT_SECURE_NO_WARNINGS
#include <Windows.h>
#include <stdio.h>
#include <string>

#include "mp_plugin.h"
#include "mp_plugin_manager.h"

#include "av_demux.h"


#define LOG_TAG "avDemux"
#include "mp_log.h"

/******************************************************************************
**    MACROS
******************************************************************************/

#define AV_IO_BUF_DEFAULT_SIZE (256* 1024)

#define AVD_COMP_NAME "av demux"
#define AVD_COMP_VER "00.01.00"

/******************************************************************************
**    VARIABLE DEFINITIONS
******************************************************************************/

/******************************************************************************
**   inner FUNCTION DEFINITIONS
******************************************************************************/
static int MAVIOReadPacket(void *opaque, uint8_t *buf, int buf_size)
{
    MPDataSource* mds = (MPDataSource*)opaque;
    
    return mds->read(buf, buf_size);
}

static int MAVIOWritePacket(void *opaque, uint8_t *buf, int buf_size)
{
    return AVERROR(EINVAL);
}

static int64_t MAVIOSeek(void *opaque, int64_t offset, int whence)
{
    MPDataSource* mds = (MPDataSource*)opaque;
    int64_t curPos = mds->tell();
    int64_t totalSize = mds->size();
    int64_t targetPos = 0;

    if (whence == AVSEEK_SIZE)
    {
        return mds->size();
    }
    int64_t new_pos = -1;
    switch (whence) {
    case SEEK_SET:
        new_pos = offset;
        break;
    case SEEK_CUR:
        new_pos += offset;
        break;
    case SEEK_END:
        new_pos = totalSize + offset;
        break;
    }

    if (new_pos >= 0)
    {
        new_pos = mds->seek(new_pos);
    }
    else
    {
        new_pos = AVERROR(EINVAL);
    }

    return new_pos;
}

static int MAVInterruptCheck(void* opaque)
{
    AvDemux* avd = (AvDemux*)opaque;

    return (avd->isExit()) ? 1 : 0;
}

/******************************************************************************
**   FUNCTION DEFINITIONS
******************************************************************************/

AvDemux::AvDemux()
    :MPDemux(AVD_COMP_NAME, AVD_COMP_VER)
{
    /** todo something */
    /* register all formats and codecs */
    av_register_all();
}

AvDemux::~AvDemux()
{
    /** todo something */

    avformat_close_input(&mCtx);

    if (mAvIo)
    {
        avio_close(mAvIo);
    }

    if (mIoBuf)
    {
        delete[] mIoBuf;
    }


    if (mSrc)
    {
        delete mSrc;
    }
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
MPError AvDemux::getContainerInfo(MPContainerInfo& info)
{
    if (mCtx)
    {
        double timeScale = av_q2d({ 1, AV_TIME_BASE });
        info.durationMs = mCtx->duration * timeScale;
        info.startTimeMs = mCtx->start_time * timeScale;
    }

    return MP_ErrorNoInit;
}

MPError AvDemux::read(MPBuffer** buffer)
{
    MPError ret = MP_ErrorBadParameter;

    if (buffer)
    {
        int vid_id = mVidId;
        int aud_id = mAudId;
        int res = 0;
        AVPacket pkt = { 0 };

        av_init_packet(&pkt);
        
        do
        {
            av_packet_unref(&pkt);
            res = av_read_frame(mCtx, &pkt);

        } while ((res == 0) && (pkt.stream_index != vid_id) && (pkt.stream_index != aud_id));

        if (0 == res)
        {
            MPBuffer* packet = new MPBuffer(pkt.size);

            if (packet)
            {
                BufferInfo& bi = packet->info();
                AVStream* st = mCtx->streams[pkt.stream_index];
                memcpy(packet->data(), pkt.data, pkt.size);
                bi.format = codecId2Fmt(st->codecpar->codec_id);
                bi.category = avmediaType2Category(st->codecpar->codec_type);
                bi.pts = pkt.pts * 1000 * st->time_base.num / st->time_base.den;
                bi.dts = pkt.dts * 1000 * st->time_base.num / st->time_base.den;
                bi.id = pkt.stream_index;
                *buffer = packet;
                ret = MP_OK;
                /** todo something */
            }
            else
            {
                LOGE("alloc packet failed\n");
                ret = MP_ErrorInsufficientResources;
            }
        }
        else
        {
            /** todo, process normal error, eof and so on */
            ret = MP_ErrorUndefined;
        }

        av_packet_unref(&pkt);

    }

    return ret;
}

MPError AvDemux::selectTrack(eMediaCategory category, int32_t idx)
{
    MPError ret = MP_ErrorInvalidStatus;
    if (mCtx)
    {
        AVFormatContext* ctx = mCtx;

        if ((0 <= idx) && (idx < ctx->nb_streams))
        {
            AVStream* st = ctx->streams[idx];

            if (st)
            {
                AVMediaType mt = category2MediaType(category);
                if (st->codecpar->codec_type == mt)
                {
                    ret = MP_OK;

                    if (eMediaCategory::Audio == category)
                    {
                        mAudId = idx;
                    }
                    else if (eMediaCategory::Video == category)
                    {
                        mVidId = idx;
                    }
                    else
                    {
                        LOGE("unspported media type: %d\n", category);
                        ret = MP_ErrorUnsupportedType;
                    }
                }
                else
                {
                    LOGE("media type mismatch, org: %d, tar:%d\n", 
                        st->codecpar->codec_type, mt);
                }
            }
            else
            {
                LOGE("the stream is null with idx(%d)\n", idx);
            }
        }
        else
        {
            LOGE("idx(%d) out of range(0 - %d)\n", idx, ctx->nb_streams);
            ret = MP_ErrorBadParameter;
        }
    }
    
    return ret;
}

MPError AvDemux::unsleectTrack(eMediaCategory category, int32_t idx)
{
    MPError ret = MP_ErrorInvalidStatus;
    if (mCtx)
    {
        AVFormatContext* ctx = mCtx;

        if ((0 <= idx) && (idx < ctx->nb_streams))
        {
            AVStream* st = ctx->streams[idx];

            if (st)
            {
                AVMediaType mt = category2MediaType(category);
                if (st->codecpar->codec_type == mt)
                {
                    ret = MP_OK;

                    if (eMediaCategory::Audio == category)
                    {
                        if (mAudId == idx)
                        {
                            mAudId = -1;
                        }
                        else
                        {
                            LOGE("target audio idx(%d) not selected\n", idx);
                            ret = MP_ErrorUndefined;
                        }
                    }
                    else if (eMediaCategory::Video == category)
                    {
                        if (mVidId == idx)
                        {
                            mVidId = -1;
                        }
                        else
                        {
                            LOGE("target video idx(%d) not selected\n", idx);
                            ret = MP_ErrorUndefined;
                        }                        
                    }
                    else
                    {
                        LOGE("unspported media type: %d\n", category);
                        ret = MP_ErrorUnsupportedType;
                    }
                }
                else
                {
                    LOGE("media type mismatch, org: %d, tar:%d\n",
                        st->codecpar->codec_type, mt);
                }
            }
            else
            {
                LOGE("the stream is null with idx(%d)\n", idx);
            }
        }
        else
        {
            LOGE("idx(%d) out of range(0 - %d)\n", idx, ctx->nb_streams);
            ret = MP_ErrorBadParameter;
        }
    }

    return ret;

}

MPError AvDemux::seek(int64_t posMs, eSeekMode mode /*= eSeekMode::PreviousSync*/)
{
    MPError ret = MP_ErrorBadParameter;

    int res = avformat_seek_file(mCtx, -1, INT64_MIN, posMs, posMs, 0);
    if (res == 0)
    {
        ret = MP_OK;
    }
    else
    {
        LOGE("could not seek to position %lld(ms), err(%d)\n", posMs, res);
    }

    return ret;
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

MPError AvDemux::initialize(const MPDemuxCreateParam* mdcp)
{
    MPError ret = MP_ErrorBadParameter;
    if (mdcp)
    {
        mAddTrack = mdcp->addTrack;
        if (mdcp->uri)
        {
            ret = initWithUrl(mdcp->uri);
        }
        else if (mdcp->source)
        {
            ret = initWithDataSource(mdcp->source);
        }
        else
        {
            LOGE("need to set uri or data source\n");
        }
    }

    return ret;
}

bool AvDemux::isExit()
{
    /** todo refine me .*/
    return false;
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

MPError AvDemux::initWithUrl(const char* uri)
{
    MPError ret = MP_ErrorInsufficientResources;
    int res = avformat_open_input(&mCtx, uri, nullptr, nullptr);

    if (res == 0)
    {
        ret = processTrackInfo();

        if (MP_OK == ret)
        {
            ret = setDefaultTrackInfo();
        }
        else
        {
            LOGE("process track info failed(%d, %s)\n", ret, MPError2Str(ret));
        }
    }
    else
    {
        LOGE("open avformat with(%s) failed(%d)\n", uri, res);
    }

    return ret;
}

MPError AvDemux::initWithDataSource(MPDataSource* src)
{
    MPError ret = MP_OK;
    mIoBuf = new uint8_t[AV_IO_BUF_DEFAULT_SIZE];

    mAvIo = avio_alloc_context(mIoBuf, AV_IO_BUF_DEFAULT_SIZE, 0, src, MAVIOReadPacket,
                               MAVIOWritePacket, MAVIOSeek);

    if (mAvIo)
    {
        AVInputFormat* fmt = NULL;

        int res = av_probe_input_buffer2(mAvIo, &fmt, NULL, NULL, 0, 0);

        if (res >= 0)
        {
            mCtx = avformat_alloc_context();

            if (mCtx)
            {
                AVFormatContext* ctx = mCtx;

                ctx->pb = mAvIo;
                ctx->interrupt_callback.opaque = src;
                ctx->interrupt_callback.callback = MAVInterruptCheck;
                // ctx->probesize = 1024 * 32 * 4;
                // ctx->max_analyze_duration = 1500;
                res = avformat_open_input(&ctx, NULL, fmt, NULL);

                if (0 == res)
                {
                    ret = processTrackInfo();

                    if (MP_OK == ret)
                    {
                        ret = setDefaultTrackInfo();
                    }
                }
                else
                {
                    LOGE("open avformat failed(0x%08x)\n", ret);
                    ret = MP_ErrorInsufficientResources;
                }
            }
            else
            {
                LOGE("aclloc avformat context failed\n");
                ret = MP_ErrorInsufficientResources;
            }
        }
        else
        {
            LOGE("probe input buffer failed\n");
            ret = MP_ErrorUndefined;
        }
    }
    else
    {
        LOGE("alloc av io failed\n");
        ret = MP_ErrorInsufficientResources;
    }

    return ret;
}

MPError AvDemux::processTrackInfo()
{
    AVFormatContext* ctx = mCtx;
    MPError ret = (ctx->nb_streams > 0) ? MP_OK : MP_ErrorUndefined;
    
    uint32_t i = 0;

    for (; i < ctx->nb_streams; i++)
    {
        AVStream* st = ctx->streams[i];

        if (st)
        {
            AVCodecParameters* codec = st->codecpar;
            enum AVMediaType codec_type = codec->codec_type;
            MPTrackInfo info;



            if (codec_type == AVMEDIA_TYPE_VIDEO)
            {
                MPVideoInfo* vt = &(info.video);

                vt->picHeight = codec->height;
                vt->picWidth = codec->width;

                vt->bufHeight = codec->height;
                vt->bufWidth = codec->width;
                vt->framerate = st->avg_frame_rate.den ? (st->avg_frame_rate.num * 1000 / st->avg_frame_rate.den) : 0;
                info.category = eMediaCategory::Video;
            }
            else if (codec_type == AVMEDIA_TYPE_AUDIO)
            {
                MPAudioInfo* at = &(info.audio);

                at->bitDepth = codec->bits_per_coded_sample;
                at->channelCount = codec->channels;
                at->samplerate = codec->sample_rate;
                info.category = eMediaCategory::Audio;
            }
            else
            {
                continue; /** not support this media type */
            }

            info.bit_rate = codec->bit_rate;
            info.format = codecId2Fmt(codec->codec_id);
            if (codec->extradata && (codec->extradata_size > 0))
            {  
                info.extraData = MAKE_SHARED_ARRAY8(codec->extradata_size);
                info.extraSize = codec->extradata_size;
                memcpy(info.extraData.get(), codec->extradata, codec->extradata_size);
            }

            info.group = 0; /** todo refine me!! */
            info.id = st->index; /** todo, refine me, id or index?? */

            //info.language
           // info.priority = 0; /** todo, refine me */

            LOGD("add track(index: %d, id: %d)\n", i, st->index);

            if (mAddTrack)
            {
                MPError result = mAddTrack(info);

                if (MP_OK != result)
                {
                    LOGE("add track to upper layer failed(%d, %s)\n", result, MPError2Str(result));
                    ret = result;
                    break;
                }
            }
        }
    }

    return ret;
}

MPError AvDemux::setDefaultTrackInfo()
{
    AVFormatContext* ctx = mCtx;
    uint32_t i = 0;
    mAudId = -1;
    mVidId = -1;

    for (; i < ctx->nb_streams; i++)
    {
        AVStream* st = ctx->streams[i];

        if (st)
        {
            AVCodecParameters* codec = st->codecpar;

            if ((codec->codec_type == AVMEDIA_TYPE_AUDIO) && (mAudId == -1))
            {
                mAudId = st->index; /** todo refine me */
            }
            else if ((codec->codec_type == AVMEDIA_TYPE_VIDEO) && (mVidId == -1))
            {
                mVidId = st->index; /** todo refine me */
            }
        }
    }

    return MP_OK;
}

eMediaFormat AvDemux::codecId2Fmt(AVCodecID id)
{
    eMediaFormat ret = eMediaFormat::Unknown;
    switch (id)
    {
    case AV_CODEC_ID_H264: ret = eMediaFormat::AVC; break;
    case AV_CODEC_ID_AAC: ret = eMediaFormat::AAC; break;
    case AV_CODEC_ID_MP3: ret = eMediaFormat::MP3; break;
    default: break;
    }

    return ret;
}

eMediaCategory AvDemux::avmediaType2Category(AVMediaType amt)
{
    eMediaCategory ret = eMediaCategory::Unknown;
    switch (amt)
    {
    case AVMEDIA_TYPE_AUDIO: ret = eMediaCategory::Audio; break;
    case AVMEDIA_TYPE_VIDEO: ret = eMediaCategory::Video; break;
    default: break;
    }

    return ret;
}

AVMediaType AvDemux::category2MediaType(eMediaCategory ct)
{
    AVMediaType ret = AVMEDIA_TYPE_UNKNOWN;
    switch (ct)
    {
    case eMediaCategory::Audio: ret = AVMEDIA_TYPE_AUDIO; break;
    case eMediaCategory::Video: ret = AVMEDIA_TYPE_VIDEO; break;
    default: break;
    }

    return ret;
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
class AvDemuxPlugin : public MPPlugin
{
public:
    AvDemuxPlugin(const std::string& name,
        const std::string& version,
        const std::string& author
    );
    virtual ~AvDemuxPlugin();

public:
    virtual MPError createComponent(const MPCCreateParam* param, MPComponent** comp) override;
};

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////
AvDemuxPlugin::AvDemuxPlugin(const std::string& name,
    const std::string& version,
    const std::string& author)
    :MPPlugin(name, version, author)
{
    /** todo something */
}

AvDemuxPlugin::~AvDemuxPlugin()
{
    /** todo something */
}

///////////////////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

MPError AvDemuxPlugin::createComponent(const MPCCreateParam* param, MPComponent** comp)
{
    MPError ret = MP_ErrorBadParameter;

    if (param && comp)
    {
        if (param->type == MP_DEMUX_COMP_TYPE)
        {
            const MPDemuxCreateParam* dp = static_cast<const MPDemuxCreateParam*>(param);

            if (dp)
            {
                AvDemux* avd = new AvDemux;

                ret = avd->initialize(dp);

                if (MP_OK == ret)
                {
                    *comp = avd;
                }
                else
                {
                    delete avd;
                }
            }
        }
        else
        {
            ret = MP_ErrorUnsupportedType;
        }
    }

    return ret;
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

/** the entry of current plugin */
MP_PLUGIN_ENTRY(AvDemuxPlugin,
    "av demux plugin",  /** name */
    "00.01.00",           /** version */
    "454727014@qq.com" /** author */
);
